"""
File Name: digits_recognizer.py
Author: Zhixuan Chang
Date: 2023-11-09
Description: recognize the input digits image and output the recognization result
"""
import torch
from torchvision.io import read_image, ImageReadMode
from torchvision.transforms import Normalize
from Net import *

import pickle
import numpy as np
import matplotlib.pyplot as plt
import argparse

norm_func = Normalize((0.1307,), (0.3081,))


def load_model() -> Net:
    model_file = 'trained_net2.pybin'
    with open(model_file, 'rb') as fp:
        mdict = pickle.load(fp)
    net = mdict['trained_net']
    return net


def load_image(img_file) -> torch.Tensor:
    # img_file = r'F:\Workspace\pycharmProjects\HandWriteDigitsRecognize\data\5.jpg'
    img = read_image(img_file, ImageReadMode.GRAY)

    # fig = plt.figure()
    # plt.axis('off')
    # plt.imshow(img.squeeze(), cmap='gray')
    # plt.show()

    return img


def digit_recognize(net, img):
    img2 = norm_func(img.float().div(255.))
    img2 = img2.view(img2.size(0), 28 * 28)

    pred = net.forward(img2).argmax(dim=1).item()

    return pred


def main():
    parser = argparse.ArgumentParser(description='recognize hand written digits')
    parser.add_argument('--image', '-i',
                        help='the image file which need to be recognized')
    args = parser.parse_args()

    net = load_model()
    img = load_image(args.image)

    pred = digit_recognize(net, img)

    plt.figure()
    plt.axis('off')
    plt.imshow(img.squeeze(), cmap='gray')
    plt.title('recognization result: ' + str(pred))
    plt.show()

    return


if __name__ == '__main__':
    main()
